package com.lf.mygpt.service;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.aliyun.dashvector.DashVectorClient;
import com.aliyun.dashvector.DashVectorCollection;
import com.aliyun.dashvector.common.DashVectorException;
import com.aliyun.dashvector.models.Doc;
import com.aliyun.dashvector.models.Vector;
import com.aliyun.dashvector.models.requests.QueryDocRequest;
import com.aliyun.dashvector.models.requests.UpsertDocRequest;
import com.aliyun.dashvector.models.responses.Response;
import com.lf.mygpt.domain.Question;
import com.lf.mygpt.service.dashscope.EmbeddingItem;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestTemplate;

import javax.annotation.PostConstruct;
import java.util.Collections;
import java.util.LinkedList;
import java.util.List;
import java.util.stream.Collectors;

/**
 * @Author lifei
 * 阿里云向量数据库相关操作
 */
@Service
public class DashVectorService {

    // 灵积大模型的 key
    @Value("${dash-scope-key}")
    private String DASH_SCOPE_KEY;
    // 向量数据库的 key
    @Value("${dash-vector-key}")
    private String DASH_VECTOR_KEY;
    private static final String COLLECTION_NAME = "test1";
    private static DashVectorClient vectorClient;

    // 灵积大模型 Embedding 的请求 Url
    private static final String TextEmbeddingUrl =
            "https://dashscope.aliyuncs.com/api/v1/services/embeddings/text-embedding/text-embedding";

    @PostConstruct
    public void initVectorClient(){
        if (vectorClient == null) {
            // 初始化向量数据库客户端
            try {
                vectorClient = new DashVectorClient(DASH_VECTOR_KEY);
            } catch (DashVectorException e) {
                throw new RuntimeException(e);
            }
        }
    }

    /**
     * 初始化问题列表到向量数据库中
     * @param questionList 问题列表
     */
    public void initQuestionToDashVector(List<Question> questionList) {
        // 阿里云文本转向量最大支持每次转25条文本，这里需要对问题列表进行分批处理
        int batchSize = 25;
        int totalSize = questionList.size(); // 总数据量
        for (int i = 0; i < totalSize; i += batchSize) {
            int endIndex = Math.min(i + batchSize, totalSize);
            // 获取当前批次的文本数据
            List<Question> batchQuestion = questionList.subList(i, endIndex);
            List<String> batchTexts = batchQuestion.stream().map(Question::getQuestion).collect(Collectors.toList());
            System.out.println(" == 处理的问题列表 ==");
            System.out.println(batchTexts);

            // 执行文本向量化操作，调用上一步构造好的 textEmbedding 方法
            List<EmbeddingItem> embeddingResultItemList = textEmbedding(batchTexts, "document");


            // 构建向量数据库插入的请求体结构
            DashVectorCollection collection = vectorClient.get(COLLECTION_NAME);
            List<Doc> docList = new LinkedList<>();
            for(int j=0; j<batchQuestion.size(); j++){
                Question question = batchQuestion.get(j);
                Vector vector = Vector.builder()
                        .value(embeddingResultItemList.get(j).getEmbedding().stream().map(Double::floatValue).collect(Collectors.toList()))
                        .build();
                Doc doc = Doc.builder()
                        .vector(vector)
                        .id(question.getQuestionId().toString())
                        .field("question", question.getQuestion())
                        .field("answer", question.getAnswer())
                        .build();
                docList.add(doc);
            }
            // 执行新增更新数据的请求
            UpsertDocRequest upsertDocRequest = UpsertDocRequest.builder()
                    .docs(docList).build();
            Response<Void> response = collection.upsert(upsertDocRequest);
            System.out.println(response);
        }
    }

    /**
     * 根据问题进行向量查询
     * @param question 用户的问题文本
     * @return 相关的初始问题知识列表
     */
    public List<Question> searchQuestion(String question) {
        // 执行用户问题向量化操作，调用上一步构造好的 textEmbedding 方法
        List<EmbeddingItem> embeddingResultItemList = textEmbedding(Collections.singletonList(question), "query");
        // 获取到我们创建好的Collection
        DashVectorCollection collection = vectorClient.get(COLLECTION_NAME);
        Vector vector = Vector.builder().value(embeddingResultItemList.get(0).getEmbedding().stream().map(Double::floatValue).collect(Collectors.toList())).build();
        QueryDocRequest queryDocRequest = QueryDocRequest.builder()
            .vector(vector)
            .topk(5)
            .build();
        Response<List<Doc>> response = collection.query(queryDocRequest);

        System.out.println(" == 向量查询结果 == ");
        System.out.println(response);
        // 将查询到的结果转换为之前构造的 Question 的格式返回给前端
        List<Question> result = new LinkedList<>();
        for (Doc doc : response.getOutput()) {
            Question question1 = new Question();
            question1.setQuestionId(Long.valueOf(doc.getId()));
            question1.setQuestion(doc.getFields().get("question").toString());
            question1.setAnswer(doc.getFields().get("answer").toString());
            result.add(question1);
        }
        return result;
    }


    /**
     * 利用灵积大模型进行文本向量化
     * @param textList  待转换的文本列表
     * @param textType  带转换的文本类型，分为 doucment 和 query，分别对应知识库文档和用户问题，阿里云文档表示对于query他们有特殊处理，使得检索效果更好
     * @return 返回文本列表对应生成的向量列表
     */
    public List<EmbeddingItem> textEmbedding(List<String> textList, String textType) {
        List<EmbeddingItem> embeddingResult = new LinkedList<>();
        // 构造请求体
        JSONObject body = new JSONObject();
        body.put("model", "text-embedding-v1");
        JSONObject texts = new JSONObject();
        JSONObject parameter = new JSONObject();
        parameter.put("text_type", textType);
        body.put("parameters", parameter);

        texts.put("texts", textList);
        body.put("input", texts);

        // 发送请求
        RestTemplate template = new RestTemplate();
        HttpEntity<String> entity = new HttpEntity<String>(body.toJSONString(), getTextEmbeddingHttpHeaders());
        ResponseEntity<JSONObject> response = template.exchange(TextEmbeddingUrl, HttpMethod.POST, entity, JSONObject.class);
        // 处理请求返回结果

        JSONArray jsonArray = response.getBody().getJSONObject("output").getJSONArray("embeddings");
        jsonArray.forEach(
                embedding -> {
                    EmbeddingItem embeddingJson = JSONObject.parseObject(JSON.toJSONString(embedding), EmbeddingItem.class);
                    embeddingResult.add(embeddingJson);
                }
        );
        return embeddingResult;
    }

    /**
     * 构造灵积大模型的请求头，请求头中需要使用 DASH_SCOPE_KEY 进行鉴权
     * @return 对应的请求投
     */
    private HttpHeaders getTextEmbeddingHttpHeaders() {
        HttpHeaders httpHeaders = new HttpHeaders();
        httpHeaders.add("Authorization", "Bearer " + DASH_SCOPE_KEY);
        httpHeaders.add("Content-Type", "application/json");
        return httpHeaders;
    }

}

